package com.yu.flow.demo;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.*;
import java.util.function.Supplier;

public class WorkflowContext {

    private final Map<String, Object> data = new ConcurrentHashMap<>();
    private final Map<String, TaskResult<?>> taskResults = new ConcurrentHashMap<>();
    private final String executionId;
    private final ExecutorService executorService;
    private final Map<String, CompletableFuture<TaskResult<?>>> taskFutures = new ConcurrentHashMap<>();

    public WorkflowContext() {
        this.executionId = UUID.randomUUID().toString();
        this.executorService = Executors.newWorkStealingPool();
    }

    public String getExecutionId() {
        return executionId;
    }

    public void addTaskResult(String taskName, TaskResult<?> result) {
        taskResults.put(taskName, result);
        if (result.getStatus() == TaskStatus.SUCCESS && result.getOutput() != null) {
            String outputKey = taskName;
            TaskConfig config = getTaskConfig(taskName);
            if (config != null && config.getOutputKey() != null) {
                outputKey = config.getOutputKey();
            }
            data.put(outputKey, result.getOutput());
        }
    }

    public TaskResult<?> getTaskResult(String taskName) {
        return taskResults.get(taskName);
    }

    public Object getData(String key) {
        return data.get(key);
    }

    public void setData(String key, Object value) {
        data.put(key, value);
    }

    public ExecutorService getExecutorService() {
        return executorService;
    }

    public void submitTask(String taskName, Supplier<TaskResult<?>> taskSupplier) {
        CompletableFuture<TaskResult<?>> future = CompletableFuture.supplyAsync(taskSupplier, executorService);
        taskFutures.put(taskName, future);
    }

    public CompletableFuture<TaskResult<?>> getTaskFuture(String taskName) {
        return taskFutures.get(taskName);
    }

    public void waitForCompletion() {
        taskFutures.values().forEach(future -> {
            try {
                future.get();
            } catch (InterruptedException | ExecutionException e) {
                Thread.currentThread().interrupt();
                throw new RuntimeException("Workflow execution interrupted", e);
            }
        });
        executorService.shutdown();
    }

    // 获取任务配置的方法（需要从工作流定义中获取，此处简化实现）
    private TaskConfig getTaskConfig(String taskName) {
        // 实际实现中需要从工作流定义中获取任务配置
        return null;
    }
}
